#!/usr/bin/python3
import sys
sys.path.insert(0, '.')
import numpy as np
import argparse
import os
import cv2
from aip import AipFace
import base64
import time
import json
import uuid
import random

""" 你的 APPID AK SK """
APP_ID = '16315779'
API_KEY = 'NqzAMufM5INE9AsSS8yw9bek'
SECRET_KEY = 'HtLXQXxdTz9gLL1hGrQ7p1y7oDcTMnN9'

# USE_FILTER_DEVICE_IDS = False
# MAX_OK_IMAGE_COUNT = None
# SKIP_IMAGE_COUNT = None
# ONLY_SEARCH = True
USE_FILTER_DEVICE_IDS = False
MAX_OK_IMAGE_COUNT = 1000
SKIP_IMAGE_COUNT = 1000
ONLY_SEARCH = False
USE_IMAGE_REGISTER = False

JSON_PATH = r'/rootfs/media/kasim/Data1/data/VideoCropFace/FaceRecJson'
IMAGE_PATH = r'/rootfs/media/kasim/Data1/data/VideoCropFace'
FACE_IMAGE_PATH = r'/rootfs/media/kasim/Data1/data/VideoCropFace'
IMAGE_LIST_FILES = [
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q0-5_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q5-10_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q10-15_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q15-20_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q20-25_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q25-30_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q30-35_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q35-40_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q40-45_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q45-50_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q50-55_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q55-60_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q60-65_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q65-70_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q70-75_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q75-80_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q80-85_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q85-90_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q90-95_filter_filter.txt',
    r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_select/face_image_q95-100_filter_filter.txt',
]
DEVICE_ID_SET = {
    '150100414a54443452067fa6d4556600',  # 104 529 19.66% 越秀星汇隽庭 7(幢/座) 1(单元) 1梯 A面
    '150100414a54443452067fa1d4096600',  #  76 391 19.44% 越秀星汇隽庭 4(幢/座) 1(单元) 2梯 A面
    '150100414a54443452064d27432e3600',  #  11  64 17.19% 大信时尚家园三期 12(幢/座) 1(单元) 1梯 A面
    '150100414a54443452064d2c43e13600',  #  33 198 16.67% 大信时尚家园二期 1(幢/座) 1(单元) 1梯 A面
    '150100414a5444345203bccb439fb500',  #  37 261 14.18% 越秀星汇隽庭 5(幢/座) 1(单元) 2梯 A面
    '150100414a54443452064d2c43ed3600',  #  70 496 14.11% 大信时尚家园三期 9(幢/座) 1(单元) 1梯 A面
    '150100414a54443452064d2742c73600',  #  52 370 14.05% 大信时尚家园二期 5(幢/座) 1(单元) 1梯 A面
    '150100414a54443452067fa6d45c6600',  #  51 399 12.78% 星光礼寓 7(幢/座) 1(单元) 1梯 A面
}


GROUP_ID = 'erised_group'
USER_ID = 'user1'


""" 读取图片 """
def get_file_content(filePath):
    with open(filePath, 'rb') as fp:
        return str(base64.b64encode(fp.read()), 'utf-8')


def conv_to_base64(image_jpg_stream):
    return str(base64.b64encode(image_jpg_stream), 'utf-8')


def get_groups(client):
    ret = client.getGroupList()
    return ret
# {
#     "group_id_list": [
#         "gid1",
#         "gid2"
#     ]
# }


def create_group(client, group_id):
    ret = get_groups(client)
    need_create_group = True
    if ret is not None:
        error_code = ret.get('error_code', -1)
        error_msg = ret.get('error_msg', '')
        if error_code != 0:
            print('get_groups error: {} {}'.format(error_code, error_msg))
        else:
            result = ret.get('result', None)
            if result is not None:
                group_id_list = result.get('group_id_list', [])
                if len(group_id_list) > 0:
                    if group_id in group_id_list:
                        need_create_group = False
    if need_create_group:
        ret = client.groupAdd(group_id)
        print(ret)


# def register(client, imagepathbase, group_id, user_id):
#
#     src_path = os.path.normpath(imagepathbase)
#     for dir_name, sub_dirs, files in os.walk(src_path):
#         dir_name = os.path.normpath(dir_name)
#         if src_path != dir_name:
#             # print(video_path, dir_name)
#             continue
#
#         for file_name in files:
#             file_ext = os.path.splitext(file_name)[1].lower()
#             if (file_ext != '.jpg'):
#                 continue
#             imagepath = os.path.join(imagepathbase, file_name)
#
#             imagestream = get_file_content(imagepath)
#
#             """ 调用人脸注册 """
#             ret = client.addUser(imagestream, "BASE64", group_id, user_id)
#             print(ret)
#             print('register', group_id, user_id)
#             time.sleep(1)


# def update(client, imagepathbase, group_id, user_id):
#
#     src_path = os.path.normpath(imagepathbase)
#     for dir_name, sub_dirs, files in os.walk(src_path):
#         dir_name = os.path.normpath(dir_name)
#         if src_path != dir_name:
#             # print(video_path, dir_name)
#             continue
#
#         count = 0
#         for file_name in files:
#             file_ext = os.path.splitext(file_name)[1].lower()
#             if (file_ext != '.jpg'):
#                 continue
#             count += 1
#             if count < 20:
#                 continue
#             imagepath = os.path.join(imagepathbase, file_name)
#
#             imagestream = get_file_content(imagepath)
#
#             """ 调用人脸更新 """
#             ret = client.updateUser(imagestream, "BASE64", group_id, user_id)
#             print(ret)
#             print('update', group_id, user_id)
#             time.sleep(1)


# def search(imagepathbase):
#
#     client = AipFace(APP_ID, API_KEY, SECRET_KEY)
#     groupId = "test"
#     for image_info in image_list:
#         userId = os.path.splitext(os.path.basename(image_info[1]))[0]
#         imagepath = os.path.join(imagepathbase, image_info[1])
#
#         image_org = cv2.imread(imagepath)
#         left, top, right, bottom = image_info[3]
#         image = image_org[top:bottom, left:right]
#         cv2.imshow('image org', image_org)
#         cv2.imshow('image', image)
#         cv2.waitKey()
#         image_jpg_stream = cv2.imencode(".jpg", image)[1]
#         imagestream = conv_to_base64(image_jpg_stream)
#
#         """ 调用人脸搜索 """
#         ret = client.search(imagestream, "BASE64", groupId)
#         print('search', userId, image_info[4])
#         print(ret)
#
#     cv2.destroyAllWindows()

register_time = 0


def register_image(client, imagestream, group_id, user_id):
    global register_time
    # 人脸注册5秒内有效，刚注册不一定马上能搜索到
    """ 调用人脸注册 """
    options = {}
    options['quality_control'] = 'NORMAL'
    ret = client.addUser(imagestream, "BASE64", group_id, user_id, options=options)
    register_time = time.time()
    return ret


def register(client, face_token, group_id, user_id):
    global register_time
    # 人脸注册5秒内有效，刚注册不一定马上能搜索到
    """ 调用人脸注册 """
    ret = client.addUser(face_token, "FACE_TOKEN", group_id, user_id)
    register_time = time.time()
    return ret
# {
#   "face_token": "2fa64a88a9d5118916f9a303782a97d3",
#   "location": {
#       "left": 117,
#       "top": 131,
#       "width": 172,
#       "height": 170,
#       "rotation": 4
#   }
# }


def update(client, face_token, group_id, user_id):
    """ 调用人脸更新 """
    ret = client.updateUser(face_token, "FACE_TOKEN", group_id, user_id)
    return ret
# {
#   "face_token": "2fa64a88a9d5118916f9a303782a97d3",
#   "location": {
#       "left": 117,
#       "top": 131,
#       "width": 172,
#       "height": 170,
#       "rotation": 4
#   }
# }


def search(client, group_id, imagestream):
    global register_time
    """ 调用人脸搜索 """
    options = {}
    options['quality_control'] = 'NORMAL'
    options['max_user_num'] = 1
    if not ONLY_SEARCH:
        cur_time = time.time()
        search_time = cur_time - register_time
        wait_time = 5 - search_time
        if wait_time > 0:
            time.sleep(wait_time)
    ret = client.search(imagestream, "BASE64", group_id, options=options)
    return ret
#   {
#     "face_token": "fid",
#     "user_list": [
#        {
#           "group_id" : "test1",
#           "user_id": "u333333",
#           "user_info": "Test User",
#           "score": 99.3
#       }
#     ]
#   }


def get_user_faces(client, group_id, user_id):
    ret = client.faceGetlist(user_id, group_id)
    return ret
# {
#   "face_list": [
#       {
#           "face_token": "fid1",
#           "ctime": "2018-01-01 00:00:00"
#       },
#       {
#           "face_token": "fid2",
#           "ctime": "2018-01-01 10:00:00"
#       }
#   ]
# }


def del_face(client, group_id, user_id, face_token):
    ret = client.faceDelete(user_id, group_id, face_token)
    return ret
# // 删除成功
# {
#   "error_code": 0,
#   "log_id": 73473737,
# }
# // 删除发生错误
# {
#   "error_code": 223106,
#   "log_id": 1382953199,
#   "error_msg": "face is not exist"
# }


def face_rec(client, group_id):
    if not os.path.exists(JSON_PATH):
        os.makedirs(JSON_PATH)
    start_time = time.time()
    for IMAGE_LIST_FILE in IMAGE_LIST_FILES:
        file_count = 0
        ok_file_count = 0
        err_file_count = 0
        base_name = os.path.splitext(os.path.basename(IMAGE_LIST_FILE))[0]
        err_file_path = os.path.join(IMAGE_PATH, '{}_rec_err.txt'.format(base_name))
        ok_file_path = os.path.join(IMAGE_PATH, '{}_rec_ok.txt'.format(base_name))
        files_set = set()
        if os.path.exists(ok_file_path):
            with open(ok_file_path, 'r') as file:
                for line in file.readlines():
                    file_name = line.split()[0]
                    files_set.add(file_name)
                    ok_file_count += 1
        if os.path.exists(err_file_path):
            with open(err_file_path, 'r') as file:
                for line in file.readlines():
                    file_name = line.split(';')[0]
                    files_set.add(file_name)
                    err_file_count += 1

        ok_file = open(ok_file_path, 'a')
        err_file = open(err_file_path, 'a')
        file_count += (ok_file_count + err_file_count)
        with open(IMAGE_LIST_FILE, 'r') as file:
            file_name_list = []
            for line in file.readlines():
                file_name = line.split()[0].strip()
                if file_name in files_set:
                    continue
                if USE_FILTER_DEVICE_IDS:
                    device_id = file_name.split('/')[-4]
                    if device_id not in DEVICE_ID_SET:
                        continue
                file_name_list.append(file_name)
            random.shuffle(file_name_list)
            count = len(file_name_list) + file_count
            if MAX_OK_IMAGE_COUNT is not None:
                if count > MAX_OK_IMAGE_COUNT:
                    count = MAX_OK_IMAGE_COUNT
            for file_name in file_name_list:
                image_file_path = os.path.join(IMAGE_PATH, file_name)
                if not os.path.exists(image_file_path):
                    continue
                file_count += 1
                if (SKIP_IMAGE_COUNT is not None) and (file_count > SKIP_IMAGE_COUNT):
                    break
                if file_count % 50 == 0:
                    ok_file.flush()
                    err_file.flush()
                    print('Proccess File: {:.2f}%, Ok {}, Err {} {:.2f}%, Time: {:.3f}s'.format(100*file_count/count, ok_file_count, err_file_count, 100*err_file_count/(err_file_count+ok_file_count), time.time()-start_time))
                imagestream = get_file_content(image_file_path)
                # user_id = os.path.splitext(file_name.replace('/', '-'))[0]
                user_id = str(uuid.uuid1()).replace('-', '')

                time.sleep(0.5)
                try:
                    ret = search(client, group_id, imagestream)
                    if ret is None:
                        err_file_count += 1
                        err_file.write(file_name + '\n')
                        continue
                    error_code = ret.get('error_code', -1)
                    error_msg = ret.get('error_msg', '')
                    if error_code not in {0, 222207}:
                        print('0 {};{};{}'.format(file_name, error_code, error_msg))
                        err_file_count += 1
                        err_file.write('{};{};{}\n'.format(file_name, error_code, error_msg))
                        continue
                    if not ONLY_SEARCH:
                        ret = ret.get('result', None)
                        if ret is not None:
                            face_token = ret.get('face_token', None)
                            if face_token is None:
                                err_file_count += 1
                                err_file.write('{};{};{}\n'.format(file_name, error_code, error_msg))
                                continue
                            user_list = ret.get('user_list', None)
                        else:
                            user_list = None
                            face_token = None
                        need_register = False
                        if user_list is None:
                            need_register = True
                        else:
                            if len(user_list) > 0:
                                user = user_list[0]
                                group_id = user["group_id"]
                                # user_info = user["user_info"]
                                score = user["score"]
                                if score < 80:
                                    need_register = True
                                else:
                                    user_id = user["user_id"]
                                    need_register = False
                                    ret = get_user_faces(client, group_id, user_id)
                                    error_code = ret.get('error_code', -1)
                                    error_msg = ret.get('error_msg', '')
                                    if error_code != 0:
                                        print('1 {};{};{}'.format(file_name, error_code, error_msg))
                                        continue
                                    result = ret.get('result', None)
                                    if result is not None:
                                        face_list = result.get("face_list", None)
                                        if face_list is not None:
                                            if len(face_list) >= 18:
                                                random.shuffle(face_list)
                                                del_face_token = face_list[0]["face_token"]
                                                del_face(client, group_id, user_id, del_face_token)
                                    # ret = update(client, face_token, group_id, user_id)
                                    if USE_IMAGE_REGISTER:
                                        ret = register_image(client, imagestream, group_id, user_id)
                                    else:
                                        ret = register(client, face_token, group_id, user_id)
                            else:
                                need_register = True

                        if need_register:
                            if (face_token is None) or USE_IMAGE_REGISTER:
                                ret = register_image(client, imagestream, group_id, user_id)
                            else:
                                ret = register(client, face_token, group_id, user_id)

                        if ret is None:
                            err_file_count += 1
                            err_file.write('{}\n'.format(file_name))
                            continue

                        error_code = ret.get('error_code', -1)
                        error_msg = ret.get('error_msg', '')
                        if error_code != 0:
                            print('2 {};{};{}'.format(file_name, error_code, error_msg))
                            err_file_count += 1
                            err_file.write('{};{};{}\n'.format(file_name, error_code, error_msg))
                            continue
                        ret = ret.get('result', None)
                        if ret is None:
                            err_file_count += 1
                            err_file.write('{};{};{}\n'.format(file_name, error_code, error_msg))
                            continue

                        face_token = ret.get('face_token', None)
                        if face_token is None:
                            err_file_count += 1
                            err_file.write('{};{};{}\n'.format(file_name, error_code, error_msg))
                            continue

                    ok_file_count += 1
                    ok_file.write('{} {}\n'.format(file_name, user_id))
                    try:
                        json_path = os.path.join(JSON_PATH, file_name).replace('.jpg', '.json')
                        json_dir = os.path.dirname(json_path)
                        if not os.path.exists(json_dir):
                            os.makedirs(json_dir)
                        with open(json_path, 'w') as json_file:
                            json.dump(ret, json_file)
                    except Exception as e:
                        print('json', e)
                except Exception as e:
                    print(e, ret)
                    err_file_count += 1
                    err_file.write(file_name + '\n')
                if ok_file_count >= count:
                    break
        ok_file.close()
        err_file.close()
        os.system('chmod a+wr {}'.format(ok_file_path))
        os.system('chmod a+wr {}'.format(err_file_path))


def main():
    random.seed(123456)
    client = AipFace(APP_ID, API_KEY, SECRET_KEY)
    client.setConnectionTimeoutInMillis(10000)
    client.setSocketTimeoutInMillis(10000)

    group_id = GROUP_ID

    create_group(client, group_id)

    face_rec(client, group_id)


if __name__ == '__main__':
    main()
